{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡1  0  0   0   S₂   -S₁⎤\n",
      "⎢                      ⎥\n",
      "⎢0  1  0  -S₂   0   S₀ ⎥\n",
      "⎢                      ⎥\n",
      "⎣0  0  1  S₁   -S₀   0 ⎦\n"
     ]
    }
   ],
   "source": [
    "from sympy import *\n",
    "\n",
    "\n",
    "# Implementation of QuaternionBase<Derived>::toRotationMatrix(void).\n",
    "# The quaternion q is given as a list [qw, qx, qy, qz].\n",
    "def QuaternionToRotationMatrix(q):\n",
    "  tx  = 2 * q[1]\n",
    "  ty  = 2 * q[2]\n",
    "  tz  = 2 * q[3]\n",
    "  twx = tx * q[0]\n",
    "  twy = ty * q[0]\n",
    "  twz = tz * q[0]\n",
    "  txx = tx * q[1]\n",
    "  txy = ty * q[1]\n",
    "  txz = tz * q[1]\n",
    "  tyy = ty * q[2]\n",
    "  tyz = tz * q[2]\n",
    "  tzz = tz * q[3]\n",
    "  return Matrix([[1 - (tyy + tzz), txy - twz, txz + twy],\n",
    "                 [txy + twz, 1 - (txx + tzz), tyz - twx],\n",
    "                 [txz - twy, tyz + twx, 1 - (txx + tyy)]])\n",
    "\n",
    "\n",
    "# Implementation of SO3Group<Scalar> expAndTheta().\n",
    "# Only implementing the first case (of very small rotation) since we take the Jacobian at zero.\n",
    "def SO3exp(omega):\n",
    "  theta = omega.norm()\n",
    "  theta_sq = theta**2\n",
    "  \n",
    "  half_theta = theta / 2\n",
    "  \n",
    "  theta_po4 = theta_sq * theta_sq\n",
    "  imag_factor = Rational(1, 2) - Rational(1, 48) * theta_sq + Rational(1, 3840) * theta_po4;\n",
    "  real_factor = 1 - Rational(1, 2) * theta_sq + Rational(1, 384) * theta_po4;\n",
    "  \n",
    "  # return SO3Group<Scalar>(Eigen::Quaternion<Scalar>(\n",
    "  #     real_factor, imag_factor * omega.x(), imag_factor * omega.y(),\n",
    "  #     imag_factor * omega.z()));\n",
    "  qw = real_factor\n",
    "  qx = imag_factor * omega[0]\n",
    "  qy = imag_factor * omega[1]\n",
    "  qz = imag_factor * omega[2]\n",
    "  \n",
    "  return QuaternionToRotationMatrix([qw, qx, qy, qz])\n",
    "\n",
    "\n",
    "# Implementation of SE3Group<Scalar> exp().\n",
    "# Only implementing the first case (of small rotation) since we take the Jacobian at zero.\n",
    "def SE3exp(tangent):\n",
    "  omega = Matrix(tangent[3:6])\n",
    "  V = SO3exp(omega)\n",
    "  rotation = V\n",
    "  translation = V * Matrix(tangent[0:3])\n",
    "  return rotation.row_join(translation)\n",
    "\n",
    "\n",
    "# Main\n",
    "init_printing(use_unicode=True)\n",
    "\n",
    "# Define the tangent vector with symbolic elements T_0 to T_5.\n",
    "# (For a matrix, use: Matrix(3, 1, lambda i,j:var('S_%d%d' % (i,j))) )\n",
    "T = Matrix(6, 1, lambda i,j:var('T_%d' % (i)))\n",
    "\n",
    "# Compute transformation matrix from tangent vector.\n",
    "T_matrix = SE3exp(T)\n",
    "\n",
    "# Define the vector current_T * src:\n",
    "S = Matrix(3, 1, lambda i,j:var('S_%d' % (i)))\n",
    "\n",
    "# Matrix-vector multiplication with homogeneous vector:\n",
    "result = T_matrix * S.col_join(Matrix([1]))\n",
    "\n",
    "# Compute Jacobian:\n",
    "# (Note: The transpose is needed for stacking the matrix columns (instead of rows) into a vector.)\n",
    "jac = result.transpose().reshape(result.rows * result.cols, 1).jacobian(T)\n",
    "\n",
    "# Take Jacobian at zero:\n",
    "jac_subs = jac.subs([(T[0], 0), (T[1], 0), (T[2], 0), (T[3], 0), (T[4], 0), (T[5], 0)])\n",
    "\n",
    "# Simplify and output:\n",
    "jac_subs_simple = simplify(jac_subs)\n",
    "pprint(jac_subs_simple)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Taking the Jacobian of these functions (sorted from inner to outer):\n",
      "<function SE3exp at 0x7fdab798a840>\n",
      "<function <lambda> at 0x7fdab7639950>\n",
      "with respect to:\n",
      "⎡T₀⎤\n",
      "⎢  ⎥\n",
      "⎢T₁⎥\n",
      "⎢  ⎥\n",
      "⎢T₂⎥\n",
      "⎢  ⎥\n",
      "⎢T₃⎥\n",
      "⎢  ⎥\n",
      "⎢T₄⎥\n",
      "⎢  ⎥\n",
      "⎣T₅⎦\n",
      "at position:\n",
      "⎡0⎤\n",
      "⎢ ⎥\n",
      "⎢0⎥\n",
      "⎢ ⎥\n",
      "⎢0⎥\n",
      "⎢ ⎥\n",
      "⎢0⎥\n",
      "⎢ ⎥\n",
      "⎢0⎥\n",
      "⎢ ⎥\n",
      "⎣0⎦\n",
      "\n",
      "Intermediate step 1, for <function SE3exp at 0x7fdab798a840>\n",
      "Position after function evaluation (function value):\n",
      "⎡1  0  0  0⎤\n",
      "⎢          ⎥\n",
      "⎢0  1  0  0⎥\n",
      "⎢          ⎥\n",
      "⎣0  0  1  0⎦\n",
      "Jacobian of this function wrt. its input only:\n",
      "⎡0  0  0  0   0   0 ⎤\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   1 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   -1  0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   -1⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  1   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   1   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  -1  0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢1  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  1  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎣0  0  1  0   0   0 ⎦\n",
      "Cumulative Jacobian wrt. the innermost parameter:\n",
      "⎡0  0  0  0   0   0 ⎤\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   1 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   -1  0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   -1⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  1   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   1   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  -1  0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢1  0  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎢0  1  0  0   0   0 ⎥\n",
      "⎢                   ⎥\n",
      "⎣0  0  1  0   0   0 ⎦\n",
      "\n",
      "Intermediate step 2, for <function <lambda> at 0x7fdab7639950>\n",
      "Position after function evaluation (function value):\n",
      "⎡S₀⎤\n",
      "⎢  ⎥\n",
      "⎢S₁⎥\n",
      "⎢  ⎥\n",
      "⎣S₂⎦\n",
      "Jacobian of this function wrt. its input only:\n",
      "⎡S₀  0   0   S₁  0   0   S₂  0   0   1  0  0⎤\n",
      "⎢                                           ⎥\n",
      "⎢0   S₀  0   0   S₁  0   0   S₂  0   0  1  0⎥\n",
      "⎢                                           ⎥\n",
      "⎣0   0   S₀  0   0   S₁  0   0   S₂  0  0  1⎦\n",
      "Cumulative Jacobian wrt. the innermost parameter:\n",
      "⎡1  0  0   0   S₂   -S₁⎤\n",
      "⎢                      ⎥\n",
      "⎢0  1  0  -S₂   0   S₀ ⎥\n",
      "⎢                      ⎥\n",
      "⎣0  0  1  S₁   -S₀   0 ⎦\n",
      "\n",
      "Final result:\n",
      "⎡1  0  0   0   S₂   -S₁⎤\n",
      "⎢                      ⎥\n",
      "⎢0  1  0  -S₂   0   S₀ ⎥\n",
      "⎢                      ⎥\n",
      "⎣0  0  1  S₁   -S₀   0 ⎦\n"
     ]
    }
   ],
   "source": [
    "# Treat the function of which we want to determine the derivative as a list of nested functions.\n",
    "# This makes it easier to compute the derivative of each part, simplify it, and concatenate the results\n",
    "# using the chain rule.\n",
    "\n",
    "### Define the function of which the Jacobian shall be taken ###\n",
    "\n",
    "# Matrix-vector multiplication with homogeneous vector:\n",
    "def MatrixVectorMultiplyHomogeneous(matrix, vector):\n",
    "  return matrix * vector.col_join(Matrix([1]))\n",
    "\n",
    "# Define the vector current_T * src:\n",
    "S = Matrix(3, 1, lambda i,j:var('S_%d' % (i)))\n",
    "\n",
    "# The list of nested functions. They will be evaluated from right to left\n",
    "# (this is to match the way they would be written in math: f(g(x)).)\n",
    "functions = [lambda matrix : MatrixVectorMultiplyHomogeneous(matrix, S), SE3exp]\n",
    "\n",
    "\n",
    "### Define the variables wrt. to take the Jacobian, and the position for evaluation ###\n",
    "\n",
    "# Chain rule:\n",
    "# d(f(g(x))) / dx = (df/dy)(g(x)) * dg/dx\n",
    "\n",
    "# Define the parameter with respect to take the Jacobian, y in the formula above:\n",
    "parameters = Matrix(6, 1, lambda i,j:var('T_%d' % (i)))\n",
    "\n",
    "# Set the position at which to take the Jacobian, g(x) in the formula above:\n",
    "parameter_values = zeros(6, 1)\n",
    "\n",
    "\n",
    "### Automatic Jacobian calculation, no need to modify anything beyond this point ###\n",
    "\n",
    "# Jacobian from previous step, dg/dx in the formula above:\n",
    "previous_jacobian = 1\n",
    "\n",
    "# TODO: Test whether this works with non-matrix functions.\n",
    "def ComputeValueAndJacobian(function, parameters, parameter_values):\n",
    "  # Evaluate the function.\n",
    "  values = function(parameter_values)\n",
    "  # Compute the Jacobian.\n",
    "  symbolic_values = function(parameters)\n",
    "  symbolic_values_vector = symbolic_values.transpose().reshape(symbolic_values.rows * symbolic_values.cols, 1)\n",
    "  parameters_vector = parameters.transpose().reshape(parameters.rows * parameters.cols, 1)\n",
    "  jacobian = symbolic_values_vector.jacobian(parameters_vector)\n",
    "  # Set in the evaluation point.\n",
    "  for row in range(0, parameters.rows):\n",
    "    for col in range(0, parameters.cols):\n",
    "      jacobian = jacobian.subs(parameters[row, col], parameter_values[row, col])\n",
    "  # Simplify the jacobian.\n",
    "  jacobian = simplify(jacobian)\n",
    "  return (values, jacobian)\n",
    "\n",
    "\n",
    "# Print info about initial state.\n",
    "print('Taking the Jacobian of these functions (sorted from inner to outer):')\n",
    "for i in range(len(functions) - 1, -1, -1):\n",
    "  print(str(functions[i]))\n",
    "print('with respect to:')\n",
    "pprint(parameters)\n",
    "print('at position:')\n",
    "pprint(parameter_values)\n",
    "print('')\n",
    "\n",
    "# Loop over all functions:\n",
    "for i in range(len(functions) - 1, -1, -1):\n",
    "  # Compute value and Jacobian of this function.\n",
    "  (values, jacobian) = ComputeValueAndJacobian(functions[i], parameters, parameter_values)\n",
    "  \n",
    "  # Update parameter_values\n",
    "  parameter_values = values\n",
    "  # Update parameters (create a new symbolic vector of the same size as parameter_values)\n",
    "  parameters = Matrix(values.rows, values.cols, lambda i,j:var('T_%d%d' % (i,j)))\n",
    "  # Concatenate this Jacobian with the previous one according to the chain rule:\n",
    "  previous_jacobian = jacobian * previous_jacobian\n",
    "  \n",
    "  # Print intermediate result\n",
    "  print('Intermediate step ' + str(len(functions) - i) + ', for ' + str(functions[i]))\n",
    "  print('Position after function evaluation (function value):')\n",
    "  pprint(parameter_values)\n",
    "  print('Jacobian of this function wrt. its input only:')\n",
    "  pprint(jacobian)\n",
    "  print('Cumulative Jacobian wrt. the innermost parameter:')\n",
    "  pprint(previous_jacobian)\n",
    "  print('')\n",
    "\n",
    "# Print final result\n",
    "print('Final result:')\n",
    "pprint(previous_jacobian)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
